import json

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.colors import ListedColormap
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import PolynomialFeatures
from sklearn.preprocessing import StandardScaler

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


class LRNL:
    """
    逻辑回归（引入多项式特征）
    """

    def __init__(self, dataset, args):
        input_dim = args.get('input_dim', 9)
        output_dim = args.get('output_dim', 1)

        # 读取数据
        df = pd.read_csv(dataset, delimiter=',', header=None)
        self.data_X = df.values[:, :input_dim]
        self.data_y = df.values[:, -output_dim]

        # 数据分割
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(self.data_X, self.data_y, random_state=666)

    # 函数定义
    def plot_decision_boundary(self, model, axis):
        x0, x1 = np.meshgrid(
            np.linspace(axis[0], axis[1], int((axis[1] - axis[0]) * 100)).reshape(-1, 1),
            np.linspace(axis[2], axis[3], int((axis[3] - axis[2]) * 100)).reshape(-1, 1),
        )
        X_new = np.c_[x0.ravel(), x1.ravel()]

        y_predict = model.predict(X_new)
        zz = y_predict.reshape(x0.shape)

        custom_cmap = ListedColormap(['#EF9A9A', '#FFF59D', '#90CAF9'])

        plt.contourf(x0, x1, zz, cmap=custom_cmap)

    def PolynomialLogisticRegression(self, degree):
        return Pipeline([
            ('poly', PolynomialFeatures(degree=degree)),
            ('std_scaler', StandardScaler()),
            ('log_reg', LogisticRegression())
        ])

    def run(self):
        # 训练模型
        log_reg = self.PolynomialLogisticRegression(degree=5)
        log_reg.fit(self.X_train, self.y_train)

        predict = log_reg.predict(self.X_test)
        all_data = np.hstack((self.X_test, self.y_test.reshape(-1, 1), predict.reshape(-1, 1)))

        cols = ["fat", "protein", "snf", "acidity", "lead", "mercury", "arsenic", "chromium", "aflatoxins", "truth", "predict"]
        raw_data = [{k: v for k, v in zip(cols, row)} for row in all_data.tolist()]

        return {
            'TrainAccuracy': log_reg.score(self.X_train, self.y_train),
            'TestAccuracy': log_reg.score(self.X_test, self.y_test),
            'Truth': list(map(int, list(self.y_test))),
            'Predict': list(map(int, list(predict))),
            'RawData': raw_data,
            'Description': '非线性逻辑回归，是一种非线性回归分析模型，常用于数据挖掘，疾病自动诊断，经济预测等领域。'
                           '它在传统逻辑回归基础上引入了多项式特征，从而赋予模型非线性拟合能力。'
                           '例如，探讨食品风险因素，并根据风险因素预测食品发生风险的概率或等级。'
                           '自变量既可以是连续的，也可以是分类的。然后通过回归分析，可以得到自变量及多项式的权重，从而可以大致了解到底哪些因素是食品风险因素。'
                           '同时根据该权值可以根据食品成分含量预测食品发生风险的可能性。'
        }


if __name__ == '__main__':
    print(LRNL('../../data/sterilizedmilk_for_classfication.csv', {}).run())

